// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: xguo

#pragma once

#include <algorithm>
#include <list>
#include <string>
#include <unordered_map>
#include <unordered_set>

#include "coin2/exchange/base/order/external.h"
#include "coin2/exchange/base/order/utils.h"

using AccountBalance = coin::proto::AccountBalance;
using CurrencyBalance = coin::proto::CurrencyBalance;
using CurrencyTransferBundle = coin::proto::CurrencyTransferBundle;
using CurrencyTransferElement = coin::proto::CurrencyTransferElement;

namespace coin2::exchange::base::order {

class AccountInfo {
 public:
  using NftTokenInfo = ::coin::proto::NftTokenInfo;
  using NftTokenInfoBundle = ::coin::proto::NftTokenInfoBundle;
  using FundingFeeElement = ::coin::proto::FundingFeeElement;
  using FundingFeeBundle = ::coin::proto::FundingFeeBundle;

  enum MsgType {
    IGNORABLE = 0,
    BALANCE_UPDATE = 1,
    POSITION_UPDATE = 2,
    TRANSFER_UPDATE = 4,
    FUNDING_FEE_UPDATE = 8,
    NFT_TOKEN_INFO_UPDATE = 16,
  };

  explicit AccountInfo(const MarketExchangeApi& mea) : mea_(mea) {}

  void SetAllZeroPosition() {
    if (!need_set_zero_position_) {
      return;
    }
    need_set_zero_position_ = false;
    for (const auto& product : product_list_) {
      if (position_map_.count(product) == 0) {
        position_map_[product] = ProductPosition();
      }
    }
  }

  void SetAllZeroBalance() {
    if (!need_set_zero_balance_) {
      return;
    }
    need_set_zero_balance_ = false;
    for (const auto& currency : currency_list_) {
      if (balance_map_.count(currency) == 0) {
        balance_map_[currency] = CurrencyBalance();
      }
    }
  }

  AccountBalance GetAccountBalance() const {
    AccountBalance account_balance;
    account_balance.set_exchange(ExchangeTypeToString(mea_.exchange));
    account_balance.set_market_type(MarketTypeToString(mea_.market));
    for (auto& ele : balance_map_) {
      if (ele.second.currency().empty()) {
        continue;
      }
      account_balance.add_each_balance()->CopyFrom(ele.second);
    }
    return account_balance;
  }

  AccountPosition GetAccountPosition() const {
    AccountPosition account_position;
    account_position.set_exchange(ExchangeTypeToString(mea_.exchange));
    account_position.set_market_type(MarketTypeToString(mea_.market));
    for (auto& ele : position_map_) {
      account_position.add_each_position()->CopyFrom(ele.second);
    }
    return account_position;
  }

  CurrencyTransferBundle GetCurrencyTransfer() const {
    CurrencyTransferBundle ccy_transfer_bundle;
    ccy_transfer_bundle.set_exchange(ExchangeTypeToString(mea_.exchange));
    ccy_transfer_bundle.set_market_type(MarketTypeToString(mea_.market));
    for (auto& ele : transfer_list_) {
      ccy_transfer_bundle.add_each_transfer()->CopyFrom(ele);
    }
    return ccy_transfer_bundle;
  }

  FundingFeeBundle GetFundingFee() const {
    FundingFeeBundle funding_fee_bundle;
    funding_fee_bundle.set_exchange(ExchangeTypeToString(mea_.exchange));
    funding_fee_bundle.set_market_type(MarketTypeToString(mea_.market));
    for (const auto& funding_fee : funding_fee_list_) {
      funding_fee_bundle.add_each_funding_fee()->CopyFrom(funding_fee);
    }
    return funding_fee_bundle;
  }

  NftTokenInfoBundle GetNftTokenInfo() const {
    NftTokenInfoBundle nft_token_info_bundle;
    nft_token_info_bundle.set_exchange(ExchangeTypeToString(mea_.exchange));
    nft_token_info_bundle.set_market_type(MarketTypeToString(mea_.market));
    for (const auto& pair : nft_token_info_map_) {
      nft_token_info_bundle.add_each_nft_token_info()->CopyFrom(pair.second);
    }
    return nft_token_info_bundle;
  }

  const auto& GetBalances() const { return balance_map_; }

  const auto& GetPositions() const { return position_map_; }

  const auto& GetNftTokenInfos() const { return nft_token_info_map_; }

  const auto& GetBalance(const std::string& currency) const { return balance_map_.at(currency); }

  const auto& GetPosition(const std::string& product) const { return position_map_.at(product); }

  const auto& GetNftTokenInfo(const std::string& product) const {
    return nft_token_info_map_.at(product);
  }

  const auto& GetBalance(const ICurrency& currency) const {
    CHECK_EQ(currency.exchange(), mea_.exchange);
    return GetBalance(currency.currency());
  }

  const auto& GetPosition(const IProduct& product) const {
    CHECK_EQ(product.exchange(), mea_.exchange);

    // For bitflyer, spot product has futures product type.
    if (product.exchange() != ExchangeType::Bitflyer) {
      CHECK_EQ(product.market(), mea_.market);
    }
    return GetPosition(product.absolute_norm());
  }

  const auto& GetNftTokenInfo(const IProduct& product) const {
    CHECK_EQ(product.exchange(), mea_.exchange);
    return GetNftTokenInfo(product.absolute_norm());
  }

  bool HasBalance(const ICurrency& currency) const {
    return balance_map_.count(currency.currency()) > 0;
  }

  bool HasPosition(const IProduct& product) const {
    return position_map_.count(product.absolute_norm()) > 0;
  }

  bool HasNftTokenInfo(const IProduct& product) const {
    return nft_token_info_map_.count(product.absolute_norm()) > 0;
  }

  void SetBalance(const std::string& currency, const CurrencyBalance& balance) {
    if (balance_map_.count(currency) > 0) {
      if (!IsClose(balance.total(), balance_map_[currency].total())) {
        bal_change_set_.insert(currency);
      }
    }

    balance_map_[currency] = balance;
  }

  void SetPosition(const std::string& product, const ProductPosition& position) {
    if (position_map_.count(product)) {
      if (!IsClose(position.net_position(), position_map_[product].net_position())) {
        pos_change_set_.insert(product);
      }
    }

    position_map_[product] = position;
  }

  void SetTransfer(const std::string& currency, const CurrencyTransferElement& transfer) {
    transfer_change_set_.insert(currency);
    transfer_list_.emplace_back(transfer);
  }

  void SetTransfer(const CurrencyTransferElement& transfer) {
    transfer_list_.emplace_back(transfer);
  }

  void SetNftTokenInfo(const std::string& symbol, const NftTokenInfo& nft_token_info) {
    nft_token_info_change_set_.insert(symbol);
    nft_token_info_map_[symbol] = nft_token_info;
  }

  void SetBalance(const ICurrency& currency, const CurrencyBalance& balance) {
    CHECK_EQ(currency.exchange(), mea_.exchange);
    SetBalance(currency.currency(), balance);
  }

  void SetPosition(const IProduct& product, const ProductPosition& position) {
    CHECK_EQ(product.exchange(), mea_.exchange);
    CHECK_EQ(product.market(), mea_.market);
    SetPosition(product.absolute_norm(), position);
  }

  void SetTransfer(const ICurrency& currency, const CurrencyTransferElement& transfer) {
    CHECK_EQ(currency.exchange(), mea_.exchange);
    SetTransfer(currency.currency(), transfer);
  }

  void SetNftTokenInfo(const IProduct& product, const NftTokenInfo& nft_token_info) {
    CHECK_EQ(product.exchange(), mea_.exchange);
    CHECK_EQ(product.market(), mea_.market);
    SetNftTokenInfo(product.absolute_norm(), nft_token_info);
  }

  void AddCurrency(const std::string& currency) {
    need_set_zero_balance_ = true;
    currency_list_.emplace(currency);
  }

  void AddProduct(const std::string& product) {
    need_set_zero_position_ = true;
    product_list_.emplace(product);
  }

  void UpdateAccountBalance(const AccountBalance& account_balance, bool all_currencies = false) {
    bal_change_set_.clear();
    for (const auto& balance : account_balance.each_balance()) {
      if (all_currencies || currency_list_.count(balance.currency()) > 0) {
        SetBalance(balance.currency(), balance);
      }
    }
    Clear();
    SetType(BALANCE_UPDATE);
  }

  void UpdateAccountPosition(const AccountPosition& account_position, bool all_products = false) {
    pos_change_set_.clear();
    for (const auto& position : account_position.each_position()) {
      if (all_products || product_list_.count(position.symbol()) > 0) {
        SetPosition(position.symbol(), position);
      }
    }
    Clear();
    SetType(POSITION_UPDATE);
  }

  void UpdateCurrencyTransfer(
      const CurrencyTransferBundle& ccy_transfer_bundle,
      bool all_currencies = false) {
    transfer_list_.clear();
    transfer_change_set_.clear();
    for (const auto& transfer : ccy_transfer_bundle.each_transfer()) {
      if (mea_.exchange == coin::proto::ExchangeType::Gdax) {
        SetTransfer(transfer);
      } else {
        if (all_currencies || currency_list_.count(transfer.currency()) > 0) {
          SetTransfer(transfer.currency(), transfer);
        }
      }
    }
    Clear();
    SetType(TRANSFER_UPDATE);
  }

  void UpdateFundingFee(const FundingFeeBundle& funding_fee_bundle) {
    funding_fee_list_.clear();
    funding_fee_change_set_.clear();
    for (const auto& funding_fee : funding_fee_bundle.each_funding_fee()) {
      if (product_list_.count(funding_fee.symbol()) == 0) {
        continue;
      }
      funding_fee_change_set_.emplace(funding_fee.symbol());
      funding_fee_list_.emplace_back(funding_fee);
    }
    Clear();
    SetType(FUNDING_FEE_UPDATE);
  }

  void UpdateNftTokenInfo(const NftTokenInfo& nft_token_info) {
    nft_token_info_change_set_.clear();
    if (product_list_.count(nft_token_info.symbol()) > 0) {
      SetNftTokenInfo(nft_token_info.symbol(), nft_token_info);
    }
    Clear();
    SetType(NFT_TOKEN_INFO_UPDATE);
  }

  void set_owner(const std::string& owner) { owner_ = owner; }

  bool is_ready() const {
    if (mea_.market == MarketType::Spot) {
      return currency_list_.size() <= balance_map_.size();
    } else if (mea_.market == MarketType::Futures) {
      return product_list_.size() <= position_map_.size();
    } else {
      CHECK_EQ(mea_.market, MarketType::Options);
      return product_list_.size() <= position_map_.size();
    }
  }

  bool is_balance_update() const { return (type_bits_ & BALANCE_UPDATE) != 0; }

  bool is_position_update() const { return (type_bits_ & POSITION_UPDATE) != 0; }

  bool is_transfer_update() const { return (type_bits_ & TRANSFER_UPDATE) != 0; }

  bool is_funding_fee_update() const { return (type_bits_ & FUNDING_FEE_UPDATE) != 0; }

  bool is_nft_token_info_update() const { return (type_bits_ & NFT_TOKEN_INFO_UPDATE) != 0; }

  bool ContainsBalanceUpdate(const std::string& currency) const {
    return bal_change_set_.count(currency) > 0;
  }

  bool ContainsPositionUpdate(const std::string& product) const {
    return pos_change_set_.count(product) > 0;
  }

  bool ContainsTransferUpdate(const std::string& currency) const {
    return transfer_change_set_.count(currency) > 0;
  }

  bool ContainsFundingFeeUpdate(const std::string& product) const {
    return funding_fee_change_set_.count(product) > 0;
  }

  bool ContainsNftTokenInfoUpdate(const std::string& product) const {
    return nft_token_info_change_set_.count(product) > 0;
  }

  const std::string& owner() const { return owner_; }
  const MarketExchangeApi& mea() const { return mea_; }
  const auto& product_list() const { return product_list_; }
  const auto& currency_list() const { return currency_list_; }

  void Clear() { type_bits_ = IGNORABLE; }

 private:
  void SetType(MsgType type) { type_bits_ |= type; }

  const MarketExchangeApi mea_;
  std::string owner_;
  std::unordered_set<std::string> pos_change_set_;
  std::unordered_set<std::string> bal_change_set_;
  std::unordered_set<std::string> transfer_change_set_;
  std::unordered_set<std::string> nft_token_info_change_set_;
  std::unordered_set<std::string> funding_fee_change_set_;

  std::unordered_set<std::string> currency_list_;
  std::unordered_set<std::string> product_list_;
  std::unordered_map<std::string, CurrencyBalance> balance_map_;
  std::unordered_map<std::string, ProductPosition> position_map_;
  std::unordered_map<std::string, NftTokenInfo> nft_token_info_map_;
  std::vector<FundingFeeElement> funding_fee_list_;
  std::list<CurrencyTransferElement> transfer_list_;
  bool need_set_zero_balance_ = false;
  bool need_set_zero_position_ = false;
  int64_t type_bits_;
};

}  // namespace coin2::exchange::base::order
